import  cv2
import os
from ultralytics import YOLO
from ultralytics.solutions import speed_estimation

def process_videocity(video_path):
    # 加载YOLO模型，用于对象检测和跟踪
    model = YOLO('yolov8n.pt')
    # 获取模型能够识别的类别名称
    names = model.names
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    # 断言视频文件是否成功打开，如果没有则抛出错误
    assert cap.isOpened(), "Error reading video file"
    # 获取视频的宽、高和帧率
    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

    # 创建输出视频文件路径
    output_video_dir = "output_videos"  # 定义输出视频文件夹
    os.makedirs(output_video_dir, exist_ok=True)  # 如果文件夹不存在，则创建它
    output_video_filename = "speed_result.mp4"
    output_video_path = os.path.join(output_video_dir, output_video_filename)

    # 创建视频写入器，用于保存处理后的视频
    video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'H264'), fps, (w, h))
    # 定义用于速度估计的水平线
    line_pts = [(0, 300), (1280, 300)]
    # 初始化速度估计器
    speed_obj = speed_estimation.SpeedEstimator(reg_pts=line_pts, names=names, view_img=True)
    # 循环读取视频帧
    while cap.isOpened():
        # 读取一帧视频
        success, im0 = cap.read()
        # 如果读取失败，则退出循环
        if not success:
            print("Video frame is empty or video processing has been successfully completed.")
            break
        # 使用模型跟踪视频中的对象
        tracks = model.track(im0, persist=True, show=False)
        # 使用速度估计器在帧上估计速度并显示结果
        im0 = speed_obj.estimate_speed(im0, tracks)
        # 将处理后的帧写入到输出视频文件
        video_writer.write(im0)
    # 释放视频捕获器资源
    cap.release()
    # 释放视频写入器资源
    video_writer.release()
    # 关闭所有OpenCV创建的窗口
    cv2.destroyAllWindows()
    # 返回输出视频文件的路径
    return output_video_path